package com.github.itdachen.gateway.WebClient;

import com.fasterxml.jackson.databind.JsonNode;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

/**
 * WebClient 的使用方法
 * 案例来源: https://mp.weixin.qq.com/s/Np9RvQcG190E1QagFPuCAw
 *
 * @author 王大宸
 * @date 2025/2/5 21:52
 */
public class WebClientTest {


    // @PostMapping(value = "/deepseek", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<Map<String, String>> chatStream(String inputPrompt) {
        Map<String, Object> message = new HashMap<>();
        message.put("role", "user");
        message.put("content", inputPrompt);

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("messages", List.of(message));
        requestBody.put("stream", true);
        requestBody.put("model", "deepseek-reason");

        WebClient webClient = WebClient.builder()
                .baseUrl("https://api.deepseek.com/v1")
                .defaultHeader("Authorization", "Bearer ${DEEPSEEK_API_KEY}")
                .build();

        return webClient.post()
                .uri("/chat/completions")
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(JsonNode.class)
                .map(this::parseDeepseekResponse)
                .takeUntil(response -> response.containsKey("finish_reason"))
                .onErrorResume(error -> Flux.just(Map.of("content", "API 调用异常：" + error.getMessage())));
    }

    private Map<String, String> parseDeepseekResponse(JsonNode response) {
        JsonNode choices = response.get("choices");
        Map<String, String> result = new HashMap<>();

        if (choices != null && choices.isArray() && !choices.isEmpty()) {
            JsonNode choice = choices.get(0);
            JsonNode delta = choice.get("delta");

            if (delta != null) {
                result.put("content",
                        Optional.ofNullable(delta.get("content"))
                                .map(JsonNode::asText)
                                .orElse(""));

                result.put("reasoning_content",
                        Optional.ofNullable(delta.get("reasoning_content"))
                                .map(JsonNode::asText)
                                .orElse(""));
            }

            Optional.ofNullable(choice.get("finish_reason"))
                    .filter(node -> !node.isNull())
                    .ifPresent(node -> result.put("finish_reason", node.asText()));
        }
        return result;
    }


}
